import torch
import torch.nn as nn
import torch.optim as optim

# 模型参数
input_size = 4
hidden_size = 8
output_size = 4
lr = 0.01
epochs = 2

# 构建带有共享参数层的多层感知机
shared_fc = nn.Linear(hidden_size, hidden_size)
MLP = nn.Sequential(nn.Linear(input_size, hidden_size), nn.ReLU(),
                    shared_fc, nn.ReLU(),
                    shared_fc, nn.ReLU(),
                    nn.Linear(hidden_size, output_size)
)

# 训练数据
X = torch.randn(1, input_size)
Y = torch.randn(1, output_size)
# 优化器
optimizer = optim.SGD(MLP.parameters(), lr=lr)
# 训练模型
for epoch in range(epochs):
    # 前向传播和计算损失
    Y_pred = MLP(X)
    loss = nn.functional.mse_loss(Y_pred, Y)
    # 反向传播和更新梯度
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # 打印每层的参数和梯度
    for name, param in MLP.named_parameters():
        print(name, param.data, param.grad)
    print('Epoch: {}, Loss: {}'.format(epoch, loss.item()))